#include "ORCL.h"

// 虽然RC是共享的，但是我们每一个上下文实例分别控制一个signal周期
void orcl_inner_group_rdma_write(struct orcl_context* orcl_ctx, int dest_inner_group_index, int source_offset, int dest_offset, int size){
    assert(orcl_ctx!=NULL);
    assert(dest_offset + size < orcl_ctx->message_buffer_size);
    
    // 向组内节点发送信息
    struct ibv_send_wr wr;
    struct ibv_send_wr *bad_send_wr;

    // 发送缓冲区
    struct ibv_sge sgl;

    // 获取完成
    struct ibv_wc wc;

    // ibv调用的返回值
    int ret;

    int opcode = IBV_WR_RDMA_WRITE;

    // 需要使用第dest_inner_group_index个rc队列来进行消息的发送
    // 初始化要发送的内容
    sgl.addr = (uint64_t)orcl_ctx->message_buffer + source_offset;
    sgl.length = size;
    sgl.lkey = orcl_ctx->rc_normal_mes_mr[dest_inner_group_index]->lkey;
    wr.opcode = (ibv_wr_opcode)opcode;
    wr.num_sge = 1;

    wr.next = NULL;
    wr.sg_list = &sgl;
    
    if(size < 170){
        wr.send_flags = IBV_SEND_INLINE;
    }else{
        wr.send_flags = 0;
    }

    // 写远程的位置
    wr.wr.rdma.remote_addr = (uint64_t)orcl_ctx->others_normal_qp[dest_inner_group_index].buf_addr + dest_offset;
    // 权限
    wr.wr.rdma.rkey = orcl_ctx->others_normal_qp[dest_inner_group_index].rkey;

    // 发送计数器自增
    orcl_ctx->rc_normal_signal_period[dest_inner_group_index] = orcl_ctx->rc_normal_signal_period[dest_inner_group_index] + 1;

    // 如果signal周期到了就signal
    if(orcl_ctx->rc_normal_signal_period[dest_inner_group_index] >= NORMAL_SIGNAL_PERIOD){
        wr.send_flags |= IBV_SEND_SIGNALED;
    }

    // 这里执行发送
    ret = ibv_post_send(orcl_ctx->rc_normal_qp[dest_inner_group_index], &wr, &bad_send_wr);

    if (ret != 0)
    {
        if (ret == ENOMEM)
        {
            printf("Send Queue is full or not enough resources to complete this operation, node_id=%d, ctx_id=%d, rc_dest=%d\n", 
                orcl_ctx->node_id, orcl_ctx->orcl_ctx_id, dest_inner_group_index);
        }
    }
    
    if(orcl_ctx->rc_normal_signal_period[dest_inner_group_index] >= NORMAL_SIGNAL_PERIOD){
        orcl_ctx->rc_normal_signal_period[dest_inner_group_index] = 0;

        // 等待一个wr完成，在多线程下，这个wr可能不是自己的
        orcl_poll_cq(orcl_ctx->rc_normal_cq[dest_inner_group_index], 1, &wc);
    }

    // 一次发送流程执行完成
}

// 在组内进行转发，转发的的发送缓冲区是forward的环形缓冲区，发送数据大小不能超过512B
void orcl_inner_group_rdma_forward(struct orcl_context* orcl_ctx, int forward_buf_index, int dest_inner_group_index, int source_offset, int dest_offset, int size){
    // 首先要保证512对齐
    assert(source_offset%512 == 0);
    // 发送的数据大小不能超过512
    assert(size<512);

    // 保证发送相关的数据结构到位
    assert(orcl_ctx != NULL && orcl_ctx->rc_forward_mr[forward_buf_index][dest_inner_group_index] != NULL && 
        orcl_ctx->rc_forward_qp[dest_inner_group_index] != NULL);
    
    // 源地址和偏移量不能让缓冲区溢出
    assert(source_offset + size < FORWARD_RC_BUF_SIZE);
    // 目标地址和偏移量也不能让缓冲区溢出
    assert(dest_offset + size < orcl_ctx->message_buffer_size);

    // 执行发送
    // 向组内节点发送信息
    struct ibv_send_wr wr;
    struct ibv_send_wr *bad_send_wr;

    // 发送缓冲区
    struct ibv_sge sgl;

    // 获取完成
    struct ibv_wc wc;

    // ibv调用的返回值
    int ret;

    int opcode = IBV_WR_RDMA_WRITE;

    // 需要使用第dest_inner_group_index个rc队列来进行消息的发送
    // 初始化要发送的内容
    // 发送的位置
    // 发送数据所在的缓冲区
    sgl.addr = (uint64_t)orcl_ctx->rc_forward_buf[forward_buf_index] + source_offset;
    sgl.length = size;

    // 发送的首位数据，魔数是不是正确
    // printf("%c\n", *(char *)orcl_ctx->rc_forward_buf[forward_buf_index] + source_offset);

    // 转发缓冲区的lkey，lkey和转发的，
    sgl.lkey = orcl_ctx->rc_forward_mr[forward_buf_index][dest_inner_group_index]->lkey;
    wr.opcode = (ibv_wr_opcode)opcode;
    wr.num_sge = 1;

    wr.next = NULL;
    wr.sg_list = &sgl;
    
    if(size < 170){
        wr.send_flags = IBV_SEND_INLINE;
    }else{
        wr.send_flags = 0;
    }

    // 规定远程要写的位置
    // 写远程的位置
    wr.wr.rdma.remote_addr = (uint64_t)orcl_ctx->others_forward_qp[dest_inner_group_index].buf_addr + dest_offset;
    // 权限
    wr.wr.rdma.rkey = orcl_ctx->others_forward_qp[dest_inner_group_index].rkey;

    // 发送计数器自增，
    orcl_ctx->rc_forward_signal_period[dest_inner_group_index] = orcl_ctx->rc_forward_signal_period[dest_inner_group_index] + 1;

    // 如果signal周期到了就signal
    if(orcl_ctx->rc_forward_signal_period[dest_inner_group_index] >= FORWARD_SIGNAL_PERIOD){
        wr.send_flags |= IBV_SEND_SIGNALED;
    }

    // 这里执行发送
    ret = ibv_post_send(orcl_ctx->rc_forward_qp[dest_inner_group_index], &wr, &bad_send_wr);

    if (ret != 0)
    {
        if (ret == ENOMEM)
        {
            printf("Send Queue is full or not enough resources to complete this operation, node_id=%d, ctx_id=%d, rc_dest=%d\n", 
                orcl_ctx->node_id, orcl_ctx->orcl_ctx_id, dest_inner_group_index);
        }
    }

    if(orcl_ctx->rc_forward_signal_period[dest_inner_group_index] >= FORWARD_SIGNAL_PERIOD){
         orcl_ctx->rc_forward_signal_period[dest_inner_group_index] = 0;

        // 等待一个wr完成，在多线程下，这个wr可能不是自己的
        orcl_poll_cq(orcl_ctx->rc_forward_cq[dest_inner_group_index], 1, &wc);
    }
}

// 向其他组的转发节点发送消息
void orcl_dci_rdma_write(struct orcl_context* orcl_ctx, int dest_group_index, int source_offset, int dest_offset, int size){
    // 不能自己发给自己
    assert(orcl_ctx->group_id != dest_group_index);
    
    // 首先向其他组发送消息
    // 首先要保证接收端512对齐，发送的数据大小不能超过512
    assert(dest_offset%512 == 0 && size<512);
    
    // 发送缓冲区和接收缓冲区不能溢出
    assert(dest_offset + size < FORWARD_DCT_BUF_SIZE && source_offset + size < orcl_ctx->message_buffer_size);

    // 和dc发送相关的信息不能少
    assert(orcl_ctx != NULL);

    // 对应转发节点的对应信息，每个组里面只有一个转发节点
    // 判断其是不是已经被正确赋值
    assert(orcl_ctx->others_dct[dest_group_index].buf_size != 0);

    int dci_index = -1;

    // 首选确定要使用的dci
    // 查看当前组有没有DCI，没有会被置为-1
    if(orcl_ctx->group_connect_dci_id[dest_group_index] == -1){
        // 这里说明当前节点没有id以供复用，我们要选一个要用的节点
        dci_index = orcl_ctx->next_dci_id;

        // 更新一下group连接的记录
        orcl_ctx->group_connect_dci_id[dest_group_index] = dci_index;

        // 更新一下dci的连接记录
        // 先将对应group使用的dci置为-1
        orcl_ctx->group_connect_dci_id[orcl_ctx->dci_connect_group_id[dci_index]] = -1;
        orcl_ctx->dci_connect_group_id[dci_index] = dest_group_index;
        
        // 更新下一个要切换连接的dci
        orcl_ctx->next_dci_id = (orcl_ctx->next_dci_id + 1) % orcl_ctx->dci_num;
    }else{
        dci_index = orcl_ctx->group_connect_dci_id[dest_group_index];
    }

    // printf("发送，使用dci:%d\n", dci_index);

    // 下面组织发送流程，使用第dci_index个dci进行发送
    struct ibv_exp_send_wr wr;
    struct ibv_exp_send_wr *bad_send_wr;

    // 发送缓冲区
    struct ibv_sge sgl;
    struct ibv_wc wc;

    int ret;

    int opcode = IBV_EXP_WR_RDMA_WRITE;

    // 查看当前数据大小
    if(size < 170){
        // 这里强制不使用dm
        // 初始化发送缓冲区，发送缓冲区是dm中的区域
        sgl.addr = (uint64_t)orcl_ctx->message_buffer + source_offset;
        sgl.length = size;
        // 当前dci的message缓冲区的lkey
        sgl.lkey = orcl_ctx->dci_mes_mr[dci_index]->lkey;

        wr.exp_opcode = (ibv_exp_wr_opcode)opcode;
        wr.num_sge = 1;
        wr.next = NULL;

        wr.sg_list = &sgl;

        wr.exp_send_flags = IBV_EXP_SEND_INLINE;    
    } else{
        // 强制使用dm
        // 在发送之前，先执行拷贝
        cpy_to_device(orcl_ctx, (uint8_t *)orcl_ctx->message_buffer, source_offset, dci_index, orcl_ctx->dm_write_offset[dci_index], size);

        sgl.addr = orcl_ctx->dm_write_offset[dci_index];
        sgl.length = size;
        // 当前dci的message缓冲区的lkey
        sgl.lkey = orcl_ctx->dci_dm_mr[dci_index]->lkey;

        wr.exp_opcode = (ibv_exp_wr_opcode)opcode;
        wr.num_sge = 1;
        wr.next = NULL;

        wr.sg_list = &sgl;

        wr.exp_send_flags = 0;
    }

    // 发往对应远程dct的缓冲区
    wr.wr.rdma.remote_addr = (uint64_t)orcl_ctx->others_dct[dest_group_index].buf_addr + dest_offset;
    // 远程缓冲区的key
    wr.wr.rdma.rkey = orcl_ctx->others_dct[dest_group_index].rkey;

    // dct独有的参数
    wr.dc.ah = orcl_ctx->other_dct_ah_attr[dci_index][dest_group_index];

    wr.dc.dct_access_key = ORCL_DCT_KEY;

    wr.dc.dct_number = orcl_ctx->others_dct[dest_group_index].dct_num;

    orcl_ctx->dci_signal_period[dci_index] = orcl_ctx->dci_signal_period[dci_index] + 1;

    if(orcl_ctx->dci_signal_period[dci_index] >= DCI_SIGNAL_PERIOD){
        // 加上signal
        wr.exp_send_flags |= IBV_EXP_SEND_SIGNALED;
    }

    // 这里执行发送
    ret = ibv_exp_post_send(orcl_ctx->dci_qp[dci_index], &wr, &bad_send_wr);

    // dm拷贝位置自增
    orcl_ctx->dm_write_offset[dci_index] = (orcl_ctx->dm_write_offset[dci_index] + 512) % DCI_DM_SIZE;

    if (ret != 0)
    {
        if (ret == ENOMEM)
        {
            printf("Send Queue is full or not enough resources to complete this operation, node_id=%d, ctx_id=%d, dci_dest=%d\n", orcl_ctx->node_id, orcl_ctx->orcl_ctx_id, dest_group_index);
        }
    }
    
    // 等待一次signal wr的完成
    if(orcl_ctx->dci_signal_period[dci_index] >= DCI_SIGNAL_PERIOD){
        orcl_ctx->dci_signal_period[dci_index] = 0;

        orcl_poll_cq(orcl_ctx->dci_cq[dci_index], 1, &wc);
    }
}